{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import *\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('LLMs-Planning/llm_planning_analysis')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from Executor import Executor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = \"LLMs-Planning/llm_planning_analysis/instances/blocksworld/mystery/generated_basic\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "# sys.path.append(\"gpt-plan-benchmark/gpt_plan_test\")\n",
    "from Executor import Executor\n",
    "# from reasoners.benchmark.bw_utils import *\n",
    "from utils import *\n",
    "from pathlib import Path\n",
    "from tarski.io import PDDLReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_problem(instance, domain):\n",
    "    reader = PDDLReader(raise_on_error=True)\n",
    "    reader.parse_domain(domain)\n",
    "    return reader.parse_instance(instance)\n",
    "def get_executor(instance, domain):\n",
    "    plan_executor = Executor(domain, instance)\n",
    "    return plan_executor\n",
    "def compute_plan(domain, instance, timeout=30):\n",
    "    fast_downward_path = os.getenv(\"FAST_DOWNWARD\")\n",
    "    # Remove > /dev/null to see the output of fast-downward\n",
    "    assert os.path.exists(f\"{fast_downward_path}/fast-downward.py\")\n",
    "    '''\n",
    "    cmd = f\"{fast_downward_path}/fast-downward.py --log-level debug --translate {domain} {instance} > /dev/null 2>&1\"\n",
    "    os.system(cmd)\n",
    "    time.sleep(0.5)\n",
    "    cmd = f\"{fast_downward_path}/fast-downward.py --log-level debug output.sas --search-options --search  \\\"astar(lmcut())\\\" > /dev/null 2>&1\" \n",
    "    os.system(cmd)\n",
    "    time.sleep(0.5)\n",
    "    '''\n",
    "    if os.path.exists(\"sas_plan\"):\n",
    "        try:\n",
    "            os.remove(\"sas_plan\")\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "\n",
    "    while not os.path.exists(\"sas_plan\"):\n",
    "        cmd = f\"timeout {timeout}s {fast_downward_path}/fast-downward.py --log-level debug {domain} {instance} --search \\\"astar(lmcut())\\\"  > /dev/null 2>&1\"\n",
    "        # print(cmd)\n",
    "        os.system(cmd)\n",
    "        time.sleep(2)\n",
    "        \n",
    "    \n",
    "    if not os.path.exists(\"sas_plan\"):\n",
    "        print(\"Plan failed\")\n",
    "        return \"\"\n",
    "    \n",
    "    return Path(\"sas_plan\").read_text()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_file = \"LLMs-Planning/llm_planning_analysis/configs/mystery_blocksworld_3.yaml\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "def read_config(config_file):\n",
    "    with open(config_file, 'r') as file:\n",
    "        data = yaml.safe_load(file)\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain_pddl = \"LLMs-Planning/llm_planning_analysis/instances/blocksworld/mystery/generated_domain.pddl\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = read_config(config_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "# re.findall(\"\\(.*?\\)\", \"(cc d), (ds)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 502/502 [17:35<00:00,  2.10s/it]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import *\n",
    "my_data = []\n",
    "for i in tqdm(os.listdir(data_path)):\n",
    "    if i.endswith(\".pddl\"):\n",
    "        data_file = os.path.join(data_path, i)\n",
    "        # print(i)\n",
    "    else:\n",
    "        print(\"Not pddl file\", i)\n",
    "        continue\n",
    "    cur_instance = data_file\n",
    "    problem = get_problem(cur_instance, domain_pddl)\n",
    "    gt_plan = compute_plan(domain_pddl, cur_instance)\n",
    "    gt_plan_text = get_plan_as_text(data)\n",
    "    my_data.append((cur_instance, gt_plan_text, len(re.findall(\"\\(.*?\\)\", gt_plan_text))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path_sup = \"LLMs-Planning/llm_planning_analysis/instances/blocksworld/mystery/generated_basic_3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 101/101 [03:32<00:00,  2.10s/it]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import *\n",
    "my_data_sup = []\n",
    "for i in tqdm(os.listdir(data_path_sup)):\n",
    "    if i.endswith(\".pddl\"):\n",
    "        data_file = os.path.join(data_path_sup, i)\n",
    "        # print(i)\n",
    "    else:\n",
    "        print(\"Not pddl file\", i)\n",
    "        continue\n",
    "    cur_instance = data_file\n",
    "    problem = get_problem(cur_instance, domain_pddl)\n",
    "    gt_plan = compute_plan(domain_pddl, cur_instance)\n",
    "    gt_plan_text = get_plan_as_text(data)\n",
    "    my_data_sup.append((cur_instance, gt_plan_text, len(re.findall(\"\\(.*?\\)\", gt_plan_text))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "502"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(my_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "101"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(my_data_sup)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "json.dump(my_data, open(\"bw_basic_mystery.json\", \"w\"))\n",
    "json.dump(my_data_sup, open(\"bw_basic_sup_mystery.json\", \"w\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "json.dump(my_data + my_data_sup, open(\"full_mystory.json\", \"w\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "441\n",
      "288\n"
     ]
    }
   ],
   "source": [
    "print(len(under_8))\n",
    "print(len(under_6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "json.dump(under_8, open(\"under_8.json\", \"w\"))\n",
    "json.dump(under_6, open(\"under_6.json\", \"w\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reproduce Haotian's 100% results\n",
    "# CUDA_VISIBLE_DEVICES=2,3 python examples/rap_blocksworld/inference.py --data_path 'examples/rap_blocksworld/data/step_4.json' --depth_limit 4 --model_dir $LLAMA2_CKPTS --lora_dir None --batch_size 1 --output_trace_in_each_iter --mem_map [16,22]\n",
    "\n",
    "# CUDA_VISIBLE_DEVICES=4,5 python examples/rap_blocksworld/inference.py --data_path 'examples/rap_blocksworld/data/under_6.json' --depth_limit 6 --model_dir $LLAMA2_CKPTS --lora_dir None --batch_size 1 --output_trace_in_each_iter --mem_map [16,22] --n_iters 20\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_data = json.load(open(\"examples/rap_blocksworld/data/full.json\", \"r\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "47\n",
      "86\n",
      "155\n",
      "153\n",
      "114\n",
      "46\n",
      "1\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "for i in range(2, 30, 2):\n",
    "    step_data = [d for d in my_data + my_data_sup if d[-1] == i]\n",
    "    if len(step_data) == 0:\n",
    "        continue\n",
    "    # json.dump(step_data, open(f\"examples/rap_blocksworld/data/new_data/step_{i}.json\", \"w\"))\n",
    "    print(len(step_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
